import random
import numpy as np

from typing import Union, List
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_moe_init_routing_quant_v2")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.tokens_shape = []
        self.topk = 1

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        if input_case.name == "x":
            self.tokens_shape = input_case.shape

        if input_case.name in ["expert_idx","rowIdx"]:
            input_case.shape[0] = self.tokens_shape[0]

            if len(input_case.shape) == 2:
                input_case.shape[1] = self.topk

        return input_case